# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/23 16:58 
# @Author  : shutao
# @FileName: stgcn_base.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import paddle

from modules.exp.base_exp import BaseExp
from modules.utils import LRScheduler


class Exp(BaseExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.base_lr = 0.05
        pass

    def get_model(self):
        from ..models import STGCN
        self.model = STGCN()
        return self.model

    def get_data_loader(self, batch_size):
        from ..data.datasets import FSDDataset

        td = FSDDataset(root_path="E:\\work\\Dataset\\FSD")
        td_loader = paddle.io.DataLoader(td, batch_size=batch_size, shuffle=True, drop_last=True)
        return td_loader, len(td_loader)

    def get_loss_function(self, **kwargs):
        return None

    def get_optimizer(self, lr: LRScheduler, parameters):
        return None

    def get_lr_scheduler(self, lr: float, iters_per_epoch, **kwargs):
        return lr

    def get_training_subject(self, data, loss_func):
        img1, img2, label = data
        img1 = paddle.cast(img1, paddle.float32)
        img2 = paddle.cast(img2, paddle.float32)
        label = paddle.cast(label, paddle.int64)

        feature_vector_1 = self.model(img1)
        feature_vector_2 = self.model(img2)
        return loss_func((feature_vector_1, feature_vector_2), label)

    def get_evaluator(self, batch_size):
        return None

    def eval(self, model, evaluator, weights):
        return None
